#ifndef __STMT_H__
#define __STMT_H__

#include <stdint.h>
#include <stdlib.h>
#include <memory.h>
#include <string.h>
#include <assert.h>
#include "ast.h"
#include "expr.h"

ast_node_t* parse_statement();

void do_print_statement(ast_node_t* node, int tab, bool s);

void print_statement(ast_node_t* node);

void print_statement_s(ast_node_t* node);

void print_expr_stmt_s(ast_node_t* node);

void print_statement_ss(ast_node_t* node, int tab);

enum {
    STATEMENT_IF,
    STATEMENT_WHILE,
    STATEMENT_FOR,
    STATEMENT_BLOCK,
    STATEMENT_EXPRESSION,
    STATEMENT_PRINT,
    STATEMENT_IMPORT,
    STATEMENT_METHOD,
    STATEMENT_CLASS,
    STATEMENT_RETURN,
    STATEMENT_GLOBAL,
    STATEMENT_BREAK,
};

typedef struct stmt_t stmt_t;
typedef struct if_stmt_t if_stmt_t;
typedef struct block_stmt_t block_stmt_t;
typedef struct expr_stmt_t expr_stmt_t;

struct stmt_t: public ast_node_t
{
    int subkind;
    location_t loc;
    stmt_t() {
        memset(this, 0, sizeof(*this));
    }
};

struct if_stmt_t: public stmt_t
{
    expr_t* cond = nullptr;
    stmt_t* then_stmt = nullptr;
    stmt_t* else_stmt = nullptr;
    if_stmt_t(): stmt_t() {
        kind = NK_IfStatement;
        subkind = STATEMENT_IF;
    }
    ~if_stmt_t() {
        free_node(cond);
        free_node(then_stmt);
        free_node(else_stmt);
        cond = nullptr;
        then_stmt = nullptr;
        else_stmt = nullptr;
    }
};

struct while_stmt_t: public stmt_t
{
    expr_t* cond = nullptr;
    stmt_t* stmt = nullptr;
    while_stmt_t(): stmt_t() {
        kind = NK_WhileStatement;
        subkind = STATEMENT_WHILE;
    }
    ~while_stmt_t()
    {
        if (cond) {
            free_node(cond);
            cond = nullptr;
        }
        if (stmt) {
            free_node(stmt);
            stmt = nullptr;
        }
    }
};

struct for_stmt_t: public stmt_t
{
    stmt_t* init = nullptr;
    stmt_t* cond = nullptr;
    stmt_t* post = nullptr;
    stmt_t* stmt = nullptr;
    for_stmt_t(): stmt_t() {
        kind = NK_ForStatement;
        subkind = STATEMENT_FOR;
    }
    ~for_stmt_t()
    {
        if (init) {
            free_node(init);
            init = nullptr;
        }
        if (cond) {
            free_node(cond);
            cond = nullptr;
        }
        if (post) {
            free_node(post);
            post = nullptr;
        }
        if (stmt) {
            free_node(stmt);
            stmt = nullptr;
        }
    }
};

typedef stmt_t* stmt_ptr_t;

DECLARE_SV(stmt_ptr_t)

struct block_stmt_t: public stmt_t
{
    ast_node_t** statements;
    block_stmt_t(): stmt_t() {
        kind = NK_CompoundStatement;
        subkind = STATEMENT_BLOCK;
        svs_init(&statements);
    }
    void push(ast_node_t* stmt);
    ~block_stmt_t() {
        for (int i = 0; i < sv_count(statements); ++i)
        {
            free_node(statements[i]);
            statements[i] = NULL;
        }
        sv_free(statements);
        statements = nullptr;
    }
};

struct expr_stmt_t: public stmt_t
{
    expr_t* expr = nullptr;
    expr_stmt_t(): stmt_t() {
        kind = NK_ExpressionStatement;
        subkind = STATEMENT_EXPRESSION;
    }
    ~expr_stmt_t() {
        free_node(expr);
        expr = nullptr;
    }
};

struct import_stmt_t: public stmt_t
{
    ast_t* ast = nullptr;
    import_stmt_t(): stmt_t() {
        kind = NK_Import;
        subkind = STATEMENT_IMPORT;
    }
    ~import_stmt_t();
};

struct return_stmt_t: public stmt_t
{
    expr_t* expr = nullptr;
    return_stmt_t(): stmt_t(), expr() {
        kind = NK_ReturnStatement;
        subkind = STATEMENT_RETURN;
    }
    ~return_stmt_t() {
        if (expr) free_node(expr);
        expr = nullptr;
    }
};

// decl actually
struct global_stmt_t: public stmt_t
{
    expr_t* expr = nullptr;
    global_stmt_t(): stmt_t(), expr() {
        kind = NK_Global;
        subkind = STATEMENT_GLOBAL;
    }
    ~global_stmt_t() {
        if (expr) free_node(expr);
        expr = nullptr;
    }
};

struct function_t;
struct param_t;
struct ptype_t;
struct stmt_t;
class any_t;

struct named_t: public ast_node_t
{
	char* name = nullptr;
	char* id = nullptr;
    named_t() {
        memset(&this->kind, 0, sizeof(ast_node_t));
    }
    ~named_t()
    {
        if (name)
            free(name);
        name = nullptr;
        if (id)
            free(id);
        id = nullptr;
    }
};

namespace stmt {

struct decl_t: public named_t
{
    expr_t* decl = nullptr;
    expr_t* init = nullptr;
    decl_t() {
        kind = NK_Declaration;
    }
    ~decl_t()
    {
        free_node(decl);
        free_node(init);
        decl = nullptr;
        init = nullptr;
    }
};

}

struct field_t: public stmt::decl_t
{
    int index;
};

struct param_t: public stmt::decl_t
{
    // bool is_var = false;
};

class env_t;

typedef void (*fn_t)(env_t* parent, env_t* env, expr_t* arg, void* bind);

struct fn_inst_t
{
    fn_t fn = nullptr;
    void* bind = nullptr;
};

struct method_t: public named_t
{
	param_t** params = nullptr;
	stmt_t** statements = nullptr;
    int param_count = 0;
    bool is_local = false;
    bool is_constructor = false;
    fn_inst_t fn;
    void* bind = nullptr;
    method_t()
    {
        kind = NK_METHOD;
    }
    void push(stmt_t* stmt)
    {
        svs_push_ptr(&statements, stmt);
    }
    void push(param_t* param) {
        svs_push_ptr(&params, param);
        ++param_count;
    }
    int count() const {
        if (!is_local)
            assert(sv_count(params) == param_count);
        return param_count;
    }
    ~method_t()
    {
        if (is_local) {
            return;
        }
        for (int i = 0; i < count(); ++i)
        {
            delete params[i];
            params[i] = nullptr;
        }
        sv_free(params);
        params = nullptr;
        for (int i = 0; i < sv_count(statements); ++i)
        {
            free_node(statements[i]);
            statements[i] = nullptr;
        }
        sv_free(statements);
        statements = nullptr;
    }
};

struct class_def_t: public named_t
{
    class_def_t* parent = nullptr;
    field_t **fields = nullptr;
    int field_count = 0;
	method_t** methods = nullptr;
    method_t** constructors = nullptr;
    method_t* default_constructor = nullptr;
    expr_t** default_fields = nullptr;
    bool is_internal_class = false;
	class_def_t() {
		kind = NK_Class;
	}
    void push(method_t* method) {
        svs_push_ptr(&methods, method);
    }
    void push(field_t* field) {
        svs_push_ptr(&fields, field);
    }
    int get_field_index(const char* name) {
        for (int i = 0; i < sv_count(fields); ++i) {
            if (strcmp(fields[i]->name, name) == 0)
                return fields[i]->index;
        }
        if (parent) {
            return parent->get_field_index(name);
        }
        return -1;
    }
    field_t* get_field(const char* name) {
        for (int i = 0; i < sv_count(fields); ++i) {
            if (strcmp(fields[i]->name, name) == 0)
                return fields[i];
        }
        if (parent) {
            return parent->get_field(name);
        }
        return nullptr;
    }
    int get_field_count() const {
        return field_count;
    }
    // int get_method(const char* name, int count) {
    //     for (int i = 0; i < sv_count(methods); ++i) {
    //         if (strcmp(methods[i]->name, name) == 0 && methods[i]->count() == count)
    //             return i;
    //     }
    //     return -1;
    // }
    method_t* get_method(const char* name, int count) {
        for (int i = 0; i < sv_count(methods); ++i) {
            if (strcmp(methods[i]->name, name) == 0 &&
                (methods[i]->count() == -1 || methods[i]->count() == count))
                return methods[i];
        }
        if (parent) {
            return parent->get_method(name, count);
        }
        return nullptr;
    }
    // int get_constructor(int count) const {
    //     for (int i = 0; i < sv_count(constructors); ++i) {
    //         if (constructors[i]->count() == count)
    //             return i;
    //     }
    //     return -1;
    // }
    method_t* get_constructor(int count) const {
        for (int i = 0; i < sv_count(constructors); ++i) {
            // printf("constructor param count: %d. %s: %d\n", i, constructors[i]->name, constructors[i]->count());
            if (constructors[i]->count() == -1 || constructors[i]->count() == count)
                return constructors[i];
        }
        return nullptr;
    }
    void to_constructor()
    {
        for (int i = 0; i < sv_count(methods); ++i) {
            if (strcmp(methods[i]->name, name) == 0 ||
                strcmp(methods[i]->name, "__init__") == 0) {
                methods[i]->is_constructor = true;
                svs_push_ptr(&constructors, methods[i]);
            }
        }
        if (sv_count(methods) == 0) {
        }
    }
	~class_def_t() {
        for (int i = 0; i < sv_count(fields); ++i)
        {
            delete fields[i];
            fields[i] = nullptr;
        }
        sv_free(fields);
        fields = nullptr;

        for (int i = 0; i < sv_count(methods); ++i)
        {
            delete methods[i];
            methods[i] = nullptr;
        }
        sv_free(methods);
        methods = nullptr;
        
        sv_free(constructors);
        constructors = nullptr;

        delete default_constructor;
        default_constructor = nullptr;

        for (int i = 0; i < sv_count(default_fields); ++i)
        {
            free_node(default_fields[i]);
            default_fields[i] = nullptr;
        }
        sv_free(default_fields);
        default_fields = nullptr;
	}
};

block_stmt_t* parse_block_stmt();

#endif  // __STMT_H__
